import os
from functools import partial
from multiprocessing.pool import Pool
from pathlib import Path
from random import sample

import cv2
import pandas as pd
from tqdm import tqdm

from constants import FACE_FORENSICS
from option import parse_args

os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"

cv2.ocl.setUseOpenCL(False)
cv2.setNumThreads(0)


def cmp_frame_name(frame_name):
    frame_name_split = frame_name.split('_')
    return int(frame_name_split[0])


def get_real_fake_pairs(root_dir, manipulated_path, mode):
    lines = open(os.path.join(root_dir, '{}.txt'.format(mode)), 'r').readlines()
    video_names = [line[:-5] for line in lines]
    pairs = []
    video_set = set(video_names)

    # for video_name in video_names:
    #     pair = (video_name, video_name)
    #     pairs.append(pair)

    for video_path in os.listdir(os.path.join(manipulated_path, 'videos')):
        video_path_split = video_path.split('_')
        assert len(video_path_split) == 2
        if video_path_split[0] not in video_set:
            continue
        pair = (video_path_split[0], video_path[:-4])
        pairs.append(pair)
    return pairs


def get_paths(vid, label, original_path, manipulated_path, crops, sample_num: int, check_mask: bool):
    ori_vid, fake_vid = vid
    ori_dir = os.path.join(original_path, crops, ori_vid)
    if label == 1:
        fake_dir = os.path.join(manipulated_path, crops, fake_vid)
    else:
        fake_dir = os.path.join(original_path, crops, ori_vid)

    # mask: both c23 and c40 depend on c23 masks
    fake_type_path = Path(manipulated_path).parent
    manipulated_path_1 = os.path.join(fake_type_path, 'c23')
    mask_dir = os.path.join(manipulated_path_1, 'diffs', fake_vid)

    data = []
    # Some unaligned videos have been removed
    # if label == 1 and not os.path.exists(fake_dir):
    #     return data
    # Some masks may have been removed
    # if label == 1 and check_mask and not os.path.exists(mask_dir):
    #     return data

    frames = os.listdir(fake_dir)
    # if len(frames) > sample_num:
    #     frames = sample(frames, sample_num)
    frames.sort(key=cmp_frame_name)
    for frame in frames:
        frame_split = str(frame).split('_')
        frame_num = int(frame_split[0])
        if frame_num % 10 != 0:
            continue
        if frame_num // 10 > sample_num:
            break
        ori_img_path = os.path.join(ori_dir, frame)
        fake_img_path = os.path.join(fake_dir, frame)
        if os.path.exists(fake_img_path) or os.path.exists(ori_img_path):
            mask_path = os.path.join(mask_dir, "{}_diff.png".format(frame[:-4]))
            if check_mask:
                if os.path.exists(mask_path):
                    # 篡改定位必须保证原始图像，伪造图像，mask同时存在
                    data.append([fake_img_path, label, ori_vid])
            else:
                # 域泛化不关心mask是否存在
                data.append([fake_img_path, label, ori_vid])
        # if label == 0 and os.path.exists(ori_img_path):
        #     data.append([ori_img_path, label, ori_vid])
        # elif label == 1 and os.path.exists(fake_img_path):
        #     mask_path = os.path.join(mask_dir, "{}_diff.png".format(frame[:-4]))
        #     if check_mask:
        #         if os.path.exists(mask_path):
        #             # 篡改定位必须保证原始图像，伪造图像，mask同时存在
        #             data.append([fake_img_path, label, ori_vid])
        #     else:
        #         # 域泛化不关心mask是否存在
        #         data.append([fake_img_path, label, ori_vid])
    return data


def save_folds(args, original_path, manipulated_path, save_path, pairs_mode, mode, sample_num=100, check_mask=True):
    data = []
    if mode == 'train_fake':
        ori_ori_pairs = set()
        ori_fake_pairs = pairs_mode
    elif mode == 'train_real':
        ori_ori_pairs = pairs_mode
        ori_fake_pairs = set()
    else:
        ori_ori_pairs = pairs_mode
        ori_fake_pairs = pairs_mode

    ori_ori_pairs = list(ori_ori_pairs)
    ori_fake_pairs = list(ori_fake_pairs)
    ori_ori_pairs.sort()
    ori_fake_pairs.sort()

    with Pool(processes=1) as p:
        # original label=0
        with tqdm(total=len(ori_ori_pairs)) as pbar:
            func = partial(get_paths, label=0, original_path=original_path, manipulated_path=manipulated_path,
                           crops=args.crops_dir, sample_num=sample_num, check_mask=check_mask)
            for v in p.imap_unordered(func, ori_ori_pairs):
                pbar.update()
                data.extend(v)
        # fake label=1
        with tqdm(total=len(ori_fake_pairs)) as pbar:
            func = partial(get_paths, label=1, original_path=original_path, manipulated_path=manipulated_path,
                           crops=args.crops_dir, sample_num=sample_num, check_mask=check_mask)
            for v in p.imap_unordered(func, ori_fake_pairs):
                pbar.update()
                data.extend(v)
    fold_data = []
    for img_path, label, ori_vid in data:
        path = Path(img_path)
        if label == 0:
            video = ori_vid
        else:
            video = path.parent.name
        file = path.name
        fold_data.append([video, file, label, ori_vid, int(file.split("_")[0])])
    # random.shuffle(fold_data)
    columns = ["video", "file", "label", "original", "frame"]
    file_name = 'data_{}_{}_{}.csv'.format(FACE_FORENSICS, args.fake_type, mode)
    save_file = os.path.join(save_path, file_name)
    pd.DataFrame(fold_data, columns=columns).to_csv(save_file, index=False)


def real_fake_split(ori_fake_pairs):
    ori_fake_pairs_real = []
    ori_fake_pairs_fake = []
    for pair in ori_fake_pairs:
        if pair[0] == pair[1]:
            ori_fake_pairs_real.append(pair)
        else:
            ori_fake_pairs_fake.append(pair)
    return ori_fake_pairs_real, ori_fake_pairs_fake


def cmp(pair):
    pair_fake_split = pair[1].split('_')
    if len(pair_fake_split) == 1:
        return pair[1]
    elif len(pair_fake_split) == 2:
        return '{}{}'.format(pair_fake_split[0], pair_fake_split[1])
    else:
        return '{}{}'.format(pair_fake_split[0], pair_fake_split[2])


def main():
    args = parse_args()
    version = args.compression_version
    save_path = '../../data/{}/{}'.format(FACE_FORENSICS, version)
    os.makedirs(save_path, exist_ok=True)
    original_path = os.path.join(args.root_dir, 'original_sequences', 'youtube', version)
    manipulated_path = os.path.join(args.root_dir, 'manipulated_sequences', args.fake_type, version)
    pairs_train = get_real_fake_pairs(args.root_dir, manipulated_path, 'train')
    pairs_val = get_real_fake_pairs(args.root_dir, manipulated_path, 'val')
    pairs_test = get_real_fake_pairs(args.root_dir, manipulated_path, 'test')
    print(len(pairs_train), len(pairs_val), len(pairs_test))

    sample_num_train = 20
    sample_num_test = 20
    sample_num_val = 10
    check_mask = True

    save_folds(args, original_path, manipulated_path, save_path, pairs_train, mode='train',
               sample_num=sample_num_train, check_mask=check_mask)
    save_folds(args, original_path, manipulated_path, save_path, pairs_test, mode='test',
               sample_num=sample_num_test, check_mask=check_mask)
    save_folds(args, original_path, manipulated_path, save_path, pairs_val, mode='val',
               sample_num=sample_num_val, check_mask=check_mask)

    ori_fake_pairs_real, ori_fake_pairs_fake = real_fake_split(pairs_train)
    save_folds(args, original_path, manipulated_path, save_path, ori_fake_pairs_fake, mode='train_real',
               sample_num=sample_num_train, check_mask=check_mask)
    save_folds(args, original_path, manipulated_path, save_path, ori_fake_pairs_fake, mode='train_fake',
               sample_num=sample_num_train, check_mask=check_mask)


if __name__ == '__main__':
    # Total: 2000
    # 'Deepfakes', 'Face2Face', 'FaceSwap', 'NeuralTextures', 'FaceShifter'
    main()
